from datasets import get_dataset_split_names
from transformers import AutoModelForSequenceClassification,AutoTokenizer,DataCollatorWithPadding,Trainer,TrainingArguments
from datasets import load_dataset
import torch
from torch.utils.data import DataLoader,Dataset,random_split
import pandas as pd
import pathlib
from torch.optim import Adam
from rich import print

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)

file_path = pathlib.Path(__file__).parent.joinpath('data').joinpath('ChnSentiCorp_htl_all.csv')
save_model_path = pathlib.Path(__file__).parent.joinpath('models').joinpath('save_dataset.pt')
save_model_checkpoint_path = pathlib.Path(__file__).parent.joinpath('models').joinpath('checkpoint')
save_model_checkpoint_path.mkdir(parents=True,exist_ok=True)

model_floder = r'D:\Models\rbt3'
# model_floder = 'hfl/rbt3'


d = load_dataset('csv',data_files=file_path.resolve().__str__(),split='train')
data_init = d.filter(lambda x:x['review'] is not None)

data_init = data_init.train_test_split(test_size=0.1)
tokenizer = AutoTokenizer.from_pretrained(model_floder)
def collate_fn_batch(batch):
    outputs = tokenizer(batch['review'],max_length=128,truncation=True)
    outputs['labels'] = torch.tensor(batch['label'])
    return outputs
data_init_all = data_init.map(collate_fn_batch,remove_columns=d.column_names)
trainset, validset = data_init_all["train"], data_init_all["test"]

collactor = DataCollatorWithPadding(tokenizer=tokenizer)

# data_train = DataLoader(trainset,batch_size=64,shuffle=True,collate_fn=collactor)
# data_test = DataLoader(validset,batch_size=64,shuffle=True,collate_fn=collactor)

train_args = TrainingArguments(
    output_dir=save_model_checkpoint_path,
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    logging_steps=10,
    evaluation_strategy="epoch",     # 评估策略
    save_strategy="epoch",           # 保存策略
    save_total_limit=3,              # 最大保存数
    learning_rate=2e-5,              # 学习率
    weight_decay=0.01,               # weight_decay
    metric_for_best_model="f1",      # 设定评估指标
)


model = AutoModelForSequenceClassification.from_pretrained(model_floder)


trainer = Trainer(model=model, 
                  args=train_args, 
                  train_dataset=trainset, 
                  eval_dataset=validset, 
                  data_collator=collactor,
                #   compute_metrics=eval_metric
                  )

trainer.train()








